package com.github.netty.util;

import javax.servlet.annotation.WebServlet;
import java.io.File;
import java.io.IOException;
import java.net.URL;
import java.net.URLDecoder;
import java.util.*;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

public class ServletAnnotationScanner {

    public static Set<Class<?>> scanPackageForWebServletAnnotations(String packageName) {
        Set<Class<?>> annotatedClasses = new HashSet<>();
        ClassLoader classLoader = Thread.currentThread().getContextClassLoader();
        assert classLoader != null;
        String path = packageName.replace('.', '/');
        Enumeration<URL> resources;
        try {
            resources = classLoader.getResources(path);
        } catch (IOException e) {
            throw new RuntimeException("Failed to get resources for package: " + packageName, e);
        }

        while (resources.hasMoreElements()) {
            URL resource = resources.nextElement();
            if ("file".equals(resource.getProtocol())) {
                scanDirectory(annotatedClasses, resource);
            } else if ("jar".equals(resource.getProtocol())) {
                scanJarFile(annotatedClasses, resource);
            }
        }

        return annotatedClasses;
    }

    private static void scanDirectory(Set<Class<?>> annotatedClasses, URL resource) {
        try {
            File directory = new File(URLDecoder.decode(resource.getFile(), "UTF-8"));
            for (File file : Objects.requireNonNull(directory.listFiles())) {
                if (file.isDirectory()) {
                    scanDirectory(annotatedClasses, file.toURI().toURL());
                } else if (file.getName().endsWith(".class")) {
                    String className = resource.getPath().replaceAll("%20", " ").replace("/", ".").replace("\\", ".");
                    className = className.substring(1, className.length() - file.getName().length() + 1) + file.getName().replace(".class", "");
                    try {
                        Class<?> clazz = Class.forName(className);
                        if (clazz.isAnnotationPresent(WebServlet.class)) {
                            annotatedClasses.add(clazz);
                        }
                    } catch (ClassNotFoundException e) {
                        throw new RuntimeException("Failed to load class: " + className, e);
                    }
                }
            }
        } catch (IOException e) {
            throw new RuntimeException("Failed to decode URL or process directory: " + resource, e);
        }
    }

    private static void scanJarFile(Set<Class<?>> annotatedClasses, URL resource) {
        try {
            JarFile jarFile = new JarFile(URLDecoder.decode(resource.getFile(), "UTF-8").substring("file:".length()).replace("%21", "!"));
            Enumeration<JarEntry> entries = jarFile.entries();
            while (entries.hasMoreElements()) {
                JarEntry entry = entries.nextElement();
                String name = entry.getName();
                if (name.endsWith(".class") && !entry.isDirectory()) {
                    String className = name.replace("/", ".").replace("\\", "").replace(".class", "");
                    try {
                        Class<?> clazz = Class.forName(className);
                        if (clazz.isAnnotationPresent(WebServlet.class)) {
                            annotatedClasses.add(clazz);
                        }
                    } catch (ClassNotFoundException e) {
                        throw new RuntimeException("Failed to load class: " + className, e);
                    }
                }
            }
        } catch (IOException e) {
            throw new RuntimeException("Failed to process JAR file: " + resource, e);
        }
    }

    public static void main(String[] args) {
        Set<Class<?>> webServletClasses = scanPackageForWebServletAnnotations("com.example.package");
        webServletClasses.forEach(System.out::println);
    }
}
